"""
# before
root
|---1.jpg
|---1.xml
|---2.jpg
|---2.xml

# after
root
|---1.jpg
|---1.txt
|---1.xml
|---2.jpg
|---2.txt
|---2.xml
|---train.txt
|---val.txt
"""

import os
import random
import xml.etree.ElementTree as ET
import glob
import cv2
import numpy as np
import math

# change_here
# ============================================================================
CLASSES = {
    "mound": 0,
}
show_path = "/project/train/show_xml2yolov5/"
#
dataroot = ["/home/data/306/", "/home/data/203/sample_m"]
global_ext = ".jpg"


# ============================================================================


class DrawLabels:
    def __init__(self, CLASSES):
        self.CLASSES = CLASSES
        self.CLASSES_COLORS = {
            self.CLASSES[j]: self.hsv2rgb(i / len(self.CLASSES) * 360, 1.0, 1.0) for i, j in
            enumerate(self.CLASSES.keys())
        }
        self.font = cv2.FONT_HERSHEY_SIMPLEX

    @staticmethod
    def hsv2rgb(h, s, v):
        h = float(h)
        s = float(s)
        v = float(v)
        h60 = h / 60.0
        h60f = math.floor(h60)
        hi = int(h60f) % 6
        f = h60 - h60f
        p = v * (1 - s)
        q = v * (1 - f * s)
        t = v * (1 - (1 - f) * s)
        r, g, b = 0, 0, 0
        if hi == 0:
            r, g, b = v, t, p
        elif hi == 1:
            r, g, b = q, v, p
        elif hi == 2:
            r, g, b = p, v, t
        elif hi == 3:
            r, g, b = p, q, v
        elif hi == 4:
            r, g, b = t, p, v
        elif hi == 5:
            r, g, b = v, p, q
        r, g, b = int(r * 255), int(g * 255), int(b * 255)
        return r, g, b

    def draw(self, img, anns):
        # anns=[
        #     [label,x_center,y_center,w,h],
        #     ...,
        #     [label,x_center,y_center,w,h],
        # ]
        img = img.copy()
        for ann in anns:
            label, x, y, w, h = ann
            img = cv2.rectangle(img, (int(x - w / 2), int(y - h / 2)), (int(x + w / 2), int(y + h / 2)),
                                self.CLASSES_COLORS[label], 2)
            img = cv2.putText(img, str(label), (int(x), int(y)), self.font, 1.2, self.CLASSES_COLORS[label], 2)
        img2 = np.zeros((img.shape[0] + 20 * (len(self.CLASSES.keys()) + 1), img.shape[1], img.shape[2])) + 128
        img2[:img.shape[0], :img.shape[1], :img.shape[2]] = img
        for i, classname in enumerate(self.CLASSES.keys()):
            img2 = cv2.putText(img2, str(self.CLASSES[classname]) + ":" + classname, (0, img.shape[0] + (i + 1) * 20),
                               self.font, 0.8, self.CLASSES_COLORS[CLASSES[classname]], 2)
        return img2


def _ltbr2xywh(box, width, height):
    x = (box[0] + box[2]) / 2.0
    y = (box[1] + box[3]) / 2.0
    w = box[2] - box[0] + 1
    h = box[3] - box[1] + 1
    x = x / width
    w = w / width
    y = y / height
    h = h / height
    return x, y, w, h


def _change_ext(path, new_ext):
    root, filename = os.path.split(path)
    name, ext = os.path.splitext(filename)
    return os.path.join(root, name + new_ext)


def _read_xml(path, height, width, classes=CLASSES):
    tree = ET.ElementTree(file=path)
    root = tree.getroot()
    anns = []
    for node in root.iter('object'):
        child = list(node)
        classname = child[0].text
        if classname not in classes:
            print(classname, "not in classes")
            continue
        label = classes[classname]
        bbox = []
        # change_here
        for atr in child[1]:
            bbox.append(float(atr.text))
        if bbox == []:
            for atr in child[4]:
                bbox.append(float(atr.text))
        if bbox == []:
            print()
        x, y, w, h = _ltbr2xywh(bbox, float(width), float(height))
        anns.append([label, x, y, w, h])
    return anns


def xml2txt(imgpaths):
    # xmlpaths = [_change_ext(i,".xml") for i in imgpaths]
    # txtpaths = [_change_ext(i,".txt") for i in imgpaths]
    valid_imgs = []
    for imgpath in imgpaths:
        # prepare
        xmlpath = _change_ext(imgpath, ".xml")
        txtpath = _change_ext(imgpath, ".txt")

        # img
        img = cv2.imread(imgpath)
        if img is None:
            print(">> imgpath: {}, cv2.imread failed".format(imgpath))
            continue
        else:
            print("     {}, cv2.imread sucess".format(imgpath))

        h, w, c = img.shape
        # print(h, w)

        # anns
        anns = _read_xml(xmlpath, h, w)
        if len(anns) > 0:
            valid_imgs.append(imgpath)
            with open(txtpath, "w") as f:
                f.write("\n".join(["{} {} {} {} {}".format(*ann) for ann in anns]))
        else:
            print(">> imgpath: {}, len(anns) == 0".format(imgpath))
    print("len(imgpaths):", len(imgpaths))
    print("len(valid_imgs):", len(valid_imgs))

    return valid_imgs


def show(imgpaths, save_path):
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    drawlabels = DrawLabels(CLASSES)
    font = cv2.FONT_HERSHEY_SIMPLEX
    for imgpath in imgpaths:
        root, filename = os.path.split(imgpath)
        img = cv2.imread(imgpath)
        height, weight, c = img.shape
        txtpath = _change_ext(imgpath, ".txt")
        with open(txtpath) as f:
            anns = f.readlines()
        anns2 = []
        for ann in anns:
            label, x, y, w, h = [float(i) for i in ann.split(" ")]
            label = int(label)
            x *= weight
            w *= weight
            y *= height
            h *= height
            anns2.append([label, x, y, w, h])
        img2 = drawlabels.draw(img, anns2)
        cv2.imwrite(save_path + "/" + filename, img2)


def main(root, ratio=0.95, save_path=None):
    assert(ratio >= 0 and ratio <= 1)
    # load imgpaths
    for i in range(len(root)):
        imgpaths = glob.glob(os.path.join(root, '*' + global_ext))
        random.shuffle(imgpaths)

        if ratio < 1:
            train_imgpaths, val_imgpaths = imgpaths[:int(len(imgpaths) * ratio)], imgpaths[int(len(imgpaths) * ratio):]         # split
            train_imgpaths, val_imgpaths = xml2txt(train_imgpaths), xml2txt(val_imgpaths)                                       # convert and save anno
        else:
            train_imgpaths = imgpaths
            train_imgpaths = xml2txt(train_imgpaths)
            val_imgpaths = train_imgpaths
        if save_path is not None:
            print(val_imgpaths)
            show(val_imgpaths, save_path)

    # save img paths
    with open(os.path.join(root, 'train.txt'), 'w') as f:
        # f.write("\n".join([os.path.split(i)[-1] for i in train_imgpaths]))
        f.write("\n".join([i for i in train_imgpaths]))
    with open(os.path.join(root, 'val.txt'), 'w') as f:
        # f.write("\n".join([os.path.split(i)[-1] for i in val_imgpaths]))
        f.write("\n".join([i for i in val_imgpaths]))


if __name__ == "__main__":
    main(dataroot, 1.0)
    # main(dataroot, 1.0, show_path)